/*
 * tree.h
 *
 *  Created on: Apr 26, 2011
 *      Author: tqlong
 */

#ifndef TREE_H_
#define TREE_H_

#include <vector>
#include "arma_vector.h"
#include "Log.h"

#define for_all_node_data(i) for (int i = 0; i < n(); i++)
#define for_all_node_child(i) for (int i = 0; i < n_child(); i++)

class Tree
{
protected:
  //typedef std::vector<vec>      dataset_type;
  typedef std::vector<int>      vint_type;
  typedef std::vector<Tree*>    vtree_type;
protected:
  const dataset_type& m_data;
  vint_type& m_idx;
  Tree *m_parent;
  int m_s, m_e;     // [m_s, m_e) interval of indices for this node
  vtree_type m_child;
public:
  Tree(const dataset_type& data)
  : m_data(data), m_parent(NULL), m_idx(*(new vint_type(m_data.size()))),
    m_s(0), m_e(m_data.size())
  {
    for_all_node_data(i) m_idx[i] = i;
  }
  Tree(Tree* parent, int s, int e)
  : m_data(parent->m_data), m_parent(parent), m_idx(parent->m_idx),
    m_s(s), m_e(e)
  {
  }
  virtual ~Tree()
  {
    if (!m_parent) {
      delete &m_idx;  // only delete indices in root
    }
    for (unsigned int i = 0; i < m_child.size(); i++) delete m_child[i];
  }

  int n() const { return m_e-m_s; }
  int n_dim() const { return m_data[0].n_elem; }
  int n_child() const { return m_child.size(); }
  int s() const { return m_s; }
  int e() const { return m_e; }

  int originalIdx(int idx) const { return m_idx[m_s+idx]; }
  int rootIdx(int idx) const { return m_s+idx; }

  const vec& getOriginal(int originalIdx) const { return m_data[originalIdx]; }
  const vec& get(int idx) const { return getOriginal(originalIdx(idx)); }

  virtual Tree* createChild(int s, int e) = 0;
  Tree& addChild(Tree* child) { m_child.push_back(child); return *this; }
  void swap(int i, int j) {
    int tmp = m_idx[m_s+i]; m_idx[m_s+i] = m_idx[m_s+j]; m_idx[m_s+j] = tmp;
  }

  void findKNN(const vec& query, std::vector<int>& knnIdx, std::vector<double>& minBound)
  {
    int K = knnIdx.size();
    for_all_node_data(i) {
      double dist = arma::norm(query-get(i), 2);
      for (int k = 0; k < K; k++) if (dist < minBound[k]) {
        knnIdx.insert(knnIdx.begin()+k, m_s+i); knnIdx.pop_back();
        minBound.insert(minBound.begin()+k, dist); minBound.pop_back();
  //        Log::m_log << "i=" << m_s+i << " (" << get(i)[0] << "," << get(i)[1] << ")" << " dist=" << dist << endl;
        break;
      }
    }
  //  for (int k = 0; k < K; k++)
  //    Log::m_log << "knn=" << knnIdx[k] << " min=" << minBound[k] << endl;
  }

};

template <typename DerivedTree>
class TreeTemplate : public Tree
{
protected:
  typedef Tree                  __Base;
  //typedef Tree::dataset_type    dataset_type;
public:
  TreeTemplate(const dataset_type& data) : __Base(data) {}
  TreeTemplate(TreeTemplate* parent, int s, int e) : __Base(parent, s, e) {}

  DerivedTree* child(int i) { return static_cast<DerivedTree*>(m_child[i]); }
  DerivedTree* self() { return static_cast<DerivedTree*>(this); }

  template <typename VISITOR> void visit(VISITOR& v) { v(self()); }  // VISITOR implements operator()(DerivedTree*)
  template <typename VISITOR> void visitRecursive(VISITOR& v, bool pre_order = true) {    // Pre-order|Post-order visiting
    if (pre_order)
      v.operator()(self());
    for_all_node_child(i) child(i)->visitRecursive(v, pre_order);
    if (!pre_order)
      v.operator()(self());
  }
  void buildTree()
  {
    self()->split();
    for_all_node_child(i) child(i)->buildTree();
  }
};

template <typename T>
class NodeStats
{
public:
  void operator()(T* node)
  {
    if (node->n_child() == 0) {
      for (int i = 0; i < node->n(); i++)
        node->stats().accumulate(node->get(i));
    }
    else {
      for (int i = 0; i < node->n_child(); i++)
        node->stats().accumulate(node->child(i)->stats());
    }
//    Log::m_log << "ms=" << node->s() << " me=" << node->e() << endl;
//    node->stats().print();
  }
};

#endif /* TREE_H_ */
